import os
from uuid import uuid4
from ultralytics import YOLO
import cv2
import json
import numpy as np

class HelmetShibie():
    def __init__(self,file,task):
        self.file = file
        self.task = task
        self.input_folder = '../../bisheqd/biyesheji/public/img/data/'
        self.output_folder = '../../bisheqd/biyesheji/public/img/results/'
        self.model = YOLO('../../xuexi1/xuexi_main/runs/detect/hardhat-detection/weights/best.pt')
        os.makedirs(self.input_folder, exist_ok=True)
        os.makedirs(self.output_folder, exist_ok=True)

    # def shibie(self):
    #     # 加载模型
    #     # 生成不重复文件名
    #     ext = os.path.splitext(self.file.name)[1]
    #     new_filename = f"{uuid4()}{ext}"
    #     input_path = os.path.join(self.input_folder, new_filename)
    #     if self.file.name.endswith(('.jpg', '.png')):
    #         file_content = self.file.read()
    #
    #         img_array = np.frombuffer(file_content, np.uint8)
    #         # 使用OpenCV从字节数组解码图像
    #         img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
    #         cv2.imwrite(input_path, img)
    #
    #         results = self.model.predict(source=img, conf=0.4, iou=0.5)
    #
    #         # 解析检测结果
    #         detection_data = []
    #         total_heads = 0
    #         total_helmets = 0
    #         confidences = []
    #
    #         for result in results:
    #             # 获取检测到的目标信息
    #             boxes = result.boxes
    #             for box in boxes:
    #                 cls_id = int(box.cls[0])  # 类别ID（例如 0=head, 1=helmet）
    #                 conf = float(box.conf[0])  # 置信度
    #                 confidences.append(conf)
    #
    #                 # 统计未戴安全帽（head）和已戴（helmet）的数量
    #                 if cls_id == 0:  # 假设类别0是未戴安全帽（head）
    #                     total_heads += 1
    #                 elif cls_id == 1:  # 假设类别1是已戴安全帽（helmet）
    #                     total_helmets += 1
    #
    #         # 总人数 = 未戴 + 已戴
    #         total_people = total_heads + total_helmets
    #         if total_heads > 0:
    #             status=1
    #         else:
    #             status=0
    #         annotated_image = results[0].plot()
    #
    #         output_path = os.path.join(self.output_folder, new_filename)
    #         cv2.imwrite(output_path, annotated_image)
    #
    #         jieguo={}
    #         jieguo["image_name"]=new_filename
    #         jieguo["people"]=total_people
    #         jieguo["nohelmet_people"]=total_heads
    #         jieguo["status"]=status
    #         # 将confidences从列表类型转成json字符串类型
    #         jieguo["confidences"]=json.dumps(confidences)
    #         return jieguo
    def _process_frame(self, frame):
        """处理单帧图像并返回标注结果和统计信息"""
        results = self.model.predict(source=frame, conf=0.6, iou=0.5)
        annotated_frame = results[0].plot()

        # 统计逻辑
        total_heads = 0
        total_helmets = 0
        confidences = []

        for box in results[0].boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            if conf > 0.5:
                confidences.append(conf)

            if cls_id == 0 and conf > 0.5:  # 未戴安全帽
                total_heads += 1
            elif cls_id == 1:  # 已戴安全帽
                total_helmets += 1

        frame_data = {
            "people": total_heads + total_helmets,
            "nohelmet_people": total_heads,
            "confidences": confidences,
            "status": 1 if total_heads > 0 else 0
        }
        return annotated_frame, frame_data

    def _process_image(self):
        """处理图片文件"""
        ext = os.path.splitext(self.file.name)[1]
        new_filename = f"{uuid4()}{ext}"
        input_path = os.path.join(self.input_folder, new_filename)
        output_path = os.path.join(self.output_folder, new_filename)

        # 保存原始图片
        img_array = np.frombuffer(self.file.read(), np.uint8)
        img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
        cv2.imwrite(input_path, img)

        # 处理并保存结果
        annotated_img, frame_data = self._process_frame(img)
        cv2.imwrite(output_path, annotated_img)

        return {
            "image_name": new_filename,
            "people": frame_data["people"],
            "nohelmet_people": frame_data["nohelmet_people"],
            "status": frame_data["status"],
            "confidences": json.dumps(frame_data["confidences"])
        }

    def _process_video(self):
        """处理视频文件"""
        ext = os.path.splitext(self.file.name)[1]
        new_filename = f"{uuid4()}{ext}"
        input_path = os.path.join(self.input_folder, new_filename)
        output_path = os.path.join(self.output_folder, new_filename)

        # 保存原始视频
        with open(input_path, 'wb') as f:
            for chunk in self.file.chunks():
                f.write(chunk)

        # 视频处理
        cap = cv2.VideoCapture(input_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fourcc = cv2.VideoWriter_fourcc(*'avc1')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

        total_heads = 0
        total_helmets = 0
        frame_details = []

        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        processed_frames = 0

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            # 处理单帧
            annotated_frame, frame_data = self._process_frame(frame)
            out.write(annotated_frame)

            processed_frames += 1
            progress = int((processed_frames / total_frames) * 100)

            # 更新进度到数据库
            self.task.progress = progress
            self.task.save()

            # 累计统计
            total_heads += frame_data["nohelmet_people"]
            total_helmets += frame_data["people"] - frame_data["nohelmet_people"]
            frame_details.append({
                "frame_confidences": frame_data["confidences"],
                "frame_status": frame_data["status"]
            })

        cap.release()
        out.release()

        return {
            "video_name": new_filename,
            "people": total_heads + total_helmets,
            "nohelmet_people": total_heads,
            "status": 1 if total_heads > 0 else 0,
            "frame_details": json.dumps(frame_details)
        }

    def shibie(self):
        """主入口方法"""
        print('名字是'+self.file.name.lower())
        if self.file.name.lower().endswith(('.jpg', '.jpeg', '.png')):
            return self._process_image()
        elif self.file.name.lower().endswith(('.mp4', '.avi', '.mov')):
            return self._process_video()
        else:
            raise ValueError("不支持的文件格式，仅支持图片(jpg/png)和视频(mp4/avi/mov)")